// src/wav-estimate.cc
// wujian@2018

// Copyright 2018 Jian Wu

// See ../../COPYING for clarification regarding multiple authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//  http://www.apache.org/licenses/LICENSE-2.0
//
// THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
// WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
// MERCHANTABLITY OR NON-INFRINGEMENT.
// See the Apache 2 License for the specific language governing permissions and
// limitations under the License.

#include "include/stft.h"

using namespace kaldi;

void EstimateSpeech(ShortTimeFTComputer &stft_computer,
                    const MatrixBase<BaseFloat> &refer_data,
                    const MatrixBase<BaseFloat> &spectrum,
                    Matrix<BaseFloat> *target_speech, bool track_volumn) {
  Matrix<BaseFloat> refer_phase, target_stft;
  // compute reference phase as target reference
  stft_computer.Compute(refer_data, NULL, NULL, &refer_phase);

  KALDI_ASSERT(SameDim(spectrum, refer_phase));
  // spectrum can be (log) magnitude/power spectrum
  stft_computer.Polar(spectrum, refer_phase, &target_stft);
  if (track_volumn) {
    BaseFloat range = refer_data.LargestAbsElem();
    stft_computer.InverseShortTimeFT(target_stft, target_speech, range);
  } else {
    stft_computer.InverseShortTimeFT(target_stft, target_speech);
  }
}

int main(int argc, char *argv[]) {
  try {
    const char *usage =
        "Estimate speech from magnitude spectrum and reference(noisy) wave\n"
        "Usage:  wav-estimate [options...] <spectrum-rspecifier> "
        "<refer-wav-rspecifier> <target-wav-wspecifier>\n"
        "   or:  wav-estimate [options...] <spectrum-rxfilename> "
        "<refer-wav-rxfilename> <target-wav-wxfilename>\n";

    ParseOptions po(usage);
    ShortTimeFTOptions stft_options;

    bool track_volumn = true;
    po.Register("track-volumn", &track_volumn,
                "If true, keep targets' volumn same as orginal wave files");

    stft_options.Register(&po);

    po.Read(argc, argv);

    if (po.NumArgs() != 3) {
      po.PrintUsage();
      exit(1);
    }

    std::string spectrum_in = po.GetArg(1), refer_in = po.GetArg(2),
                target_out = po.GetArg(3);

    bool spectrum_is_rspecifier =
             (ClassifyRspecifier(spectrum_in, NULL, NULL) != kNoRspecifier),
         refer_is_rspecifier =
             (ClassifyRspecifier(refer_in, NULL, NULL) != kNoRspecifier),
         target_is_wspecifier = (ClassifyWspecifier(target_out, NULL, NULL,
                                                    NULL) != kNoWspecifier);

    if (spectrum_is_rspecifier != target_is_wspecifier)
      KALDI_ERR << "Cannot mix archives with regular files";

    if (spectrum_is_rspecifier != refer_is_rspecifier)
      KALDI_ERR << "Configure with spectrum and reference wave must keep same";

    ShortTimeFTComputer stft_computer(stft_options);

    if (spectrum_is_rspecifier) {
      SequentialBaseFloatMatrixReader spectrum_reader(spectrum_in);
      RandomAccessTableReader<WaveHolder> refer_reader(refer_in);
      TableWriter<WaveHolder> wav_writer(target_out);

      int num_utts = 0, num_no_tgt_utts = 0, num_done = 0;
      for (; !spectrum_reader.Done(); spectrum_reader.Next()) {
        std::string utt_key = spectrum_reader.Key();
        num_utts += 1;

        if (!refer_reader.HasKey(utt_key)) {
          KALDI_WARN << utt_key << ", missing target masks";
          num_no_tgt_utts++;
          continue;
        }

        Matrix<BaseFloat> &spectrum = spectrum_reader.Value();
        const WaveData &refer_data = refer_reader.Value(utt_key);
        BaseFloat target_freq = refer_data.SampFreq();

        KALDI_ASSERT(refer_data.Data().NumRows() == 1);

        Matrix<BaseFloat> target_speech;
        EstimateSpeech(stft_computer, refer_data.Data(), spectrum,
                       &target_speech, track_volumn);

        WaveData target_data(target_freq, target_speech);
        wav_writer.Write(utt_key, target_data);
        num_done++;

        if (num_done % 100 == 0)
          KALDI_LOG << "Processed " << num_utts << " utterances";
        KALDI_VLOG(2) << "Estimate target for utterance " << utt_key;
      }
      KALDI_LOG << "Done " << num_done << " utterances out of " << num_utts
                << ", " << num_no_tgt_utts << " missing targets masks";
      return num_done == 0 ? 1 : 0;
    } else {
      Matrix<BaseFloat> spectrum;
      ReadKaldiObject(spectrum_in, &spectrum);

      bool binary;
      Input ki(refer_in, &binary);
      WaveData refer_data;
      refer_data.Read(ki.Stream());
      BaseFloat target_freq = refer_data.SampFreq();
      KALDI_ASSERT(refer_data.Data().NumRows() == 1);

      Matrix<BaseFloat> target_speech;
      EstimateSpeech(stft_computer, refer_data.Data(), spectrum, &target_speech,
                     track_volumn);

      Output ko(target_out, binary, false);
      WaveData target_data(target_freq, target_speech);
      target_data.Write(ko.Stream());

      KALDI_LOG << "Done processed " << spectrum_in;
    }

  } catch (const std::exception &e) {
    std::cerr << e.what();
    return -1;
  }
  return 0;
}
